library(xtable)
library(dplyr)
library(ggplot2)
library(scales)
library(extrafont)

rm(list = ls())
setwd("")

# Constants
SIGMA2 <- 1
VSTAR <- .002
A <- .25
T <- 3
F <- 2

# Survey cost function
s <- function(mode, multiple.measures){
  stopifnot(mode == 'O' | mode == 'T')
  stopifnot(multiple.measures == 0 | multiple.measures == 1)
  if(mode == 'T' & multiple.measures == 1) return(10)
  else return(5)
}

# rho^2 map
rho2 <- function(mode, multiple.measures){
  if(mode == 'T' & multiple.measures == 0) return(.16)
  if(mode == 'O' & multiple.measures == 0) return(.5)
  if(mode == 'T' & multiple.measures == 1) return(.33)
  if(mode == 'O' & multiple.measures == 1) return(.75)
}

# Survey response rate maps
r1 <- .07
r2 <- function(mode){
  if(mode == 'T') return(.35)
  if(mode == 'O') return(.75)
  stop('Mode must be T or O.')
}

# Formulas
placebo0pre0 <- function(mode, multiple.measures) 4 * (1/A^2) * (SIGMA2 / VSTAR) * (s(mode, multiple.measures)*F + T / (2 * r1))
placebo0pre1 <- function(mode, multiple.measures) 4 * (1/A^2) * (1 - rho2(mode, multiple.measures)) *
  (SIGMA2 / VSTAR) * (s(mode, multiple.measures)*F + T / (2 * r2(mode)) + s(mode, multiple.measures) / r2(mode))
placebo1pre0 <- function(mode, multiple.measures) 4 * (SIGMA2 / VSTAR) * (s(mode, multiple.measures)*F + T / (A * r1))
placebo1pre1 <- function(mode, multiple.measures) 4 * (1 - rho2(mode, multiple.measures)) *
  (SIGMA2 / VSTAR) * (s(mode, multiple.measures)*F + T / (A * r2(mode)) + s(mode, multiple.measures) / (A * r2(mode)))

get.design.cost <- function(pre, placebo, multiple.measures, mode){
  if(placebo == 0 & pre == 0) return(placebo0pre0(mode, multiple.measures))
  if(placebo == 1 & pre == 0) return(placebo1pre0(mode, multiple.measures))
  if(placebo == 0 & pre == 1) return(placebo0pre1(mode, multiple.measures))
  if(placebo == 1 & pre == 1) return(placebo1pre1(mode, multiple.measures))
  stop('Placebo and pre must be 0 or 1.')
}

get.design.cost.wrapper <- function(row) get.design.cost(row[1], row[2], row[3], row[4])

all.possible.designs <- expand.grid(pre = c(1,0),
                                    placebo = c(1,0),
                                    multiple.measures = c(1,0),
                                    mode = c('O', 'T'))
all.possible.designs$mode <- as.character(all.possible.designs$mode)

all.possible.designs$cost <- apply(all.possible.designs, 1, get.design.cost.wrapper)

# Aesthetic stuff
all.possible.designs$prettyname <- paste0('Pre=',all.possible.designs$pre,
                                          ', Placebo=', all.possible.designs$placebo,
                                          ', Multiple Measures=',all.possible.designs$multiple.measures,
                                          ', Mode=',all.possible.designs$mode)
all.possible.designs[all.possible.designs == 1] <- 'checkmark'
all.possible.designs[all.possible.designs == 0] <- 'No'
all.possible.designs[all.possible.designs == 'O'] <- 'Online'
all.possible.designs[all.possible.designs == 'T'] <- 'Telephone'

print.table <- function(apd){
  apd <- select(apd, -prettyname)
  apd$cost<- paste0("$",
                    format(
                      round(apd$cost, digits = -3
                      ),
                      big.mark=","),
                    sep="")
  names(apd) <- c('Pre?', 'Placebo?', 'Multiple Measures?', 'Mode?', 'ccdot')
  print(xtable(apd), include.rownames = FALSE)
}

print.table(all.possible.designs)
print.table(all.possible.designs[order(all.possible.designs$cost),]) # sort by cost

# Bar graph
all.possible.designs <- all.possible.designs[order(all.possible.designs$cost),]
all.possible.designs <- select(all.possible.designs, cost, prettyname)
all.possible.designs

#Other Possible Designs
all.possible.designs$label <- 'Other Possible Designs'
#Traditional Design
all.possible.designs$label[all.possible.designs$prettyname == 'Pre=0, Placebo=0, Multiple Measures=0, Mode=T'] <- 'Traditional Design (No Studied Practices)'
#Proposed Design
all.possible.designs$label[all.possible.designs$prettyname == 'Pre=1, Placebo=1, Multiple Measures=1, Mode=O'] <- 'All Four Practices'
#Designs Previously Used in Literature (other than traditional)
all.possible.designs$label[all.possible.designs$prettyname == 'Pre=0, Placebo=0, Multiple Measures=1, Mode=T'] <- 'Designs Previously Used'
all.possible.designs$label[all.possible.designs$prettyname == 'Pre=1, Placebo=1, Multiple Measures=0, Mode=T'] <- 'Designs Previously Used'
all.possible.designs$label[all.possible.designs$prettyname == 'Pre=1, Placebo=0, Multiple Measures=0, Mode=T'] <- 'Designs Previously Used'
all.possible.designs$label[all.possible.designs$prettyname == 'Pre=0, Placebo=0, Multiple Measures=1, Mode=T'] <- 'Designs Previously Used'
all.possible.designs$label[all.possible.designs$prettyname == 'Pre=1, Placebo=0, Multiple Measures=1, Mode=T'] <- 'Designs Previously Used'

palette <- c('firebrick1', 'black', 'gray77', 'dodgerblue3')
g <- ggplot(all.possible.designs,
       aes(x = factor(prettyname, levels = prettyname), y = cost)) +
  geom_bar(aes(fill = factor(label)), stat = 'identity') +
  ylab('Variable Cost') + xlab('Design') +
  scale_y_continuous(labels = scales::dollar) +
  coord_flip() +
  theme(text = element_text(family = 'Verdana', size = 10), legend.position="bottom") + 
  guides(fill=guide_legend(title=NULL, nrow=2)) + 
  scale_fill_manual(values=palette) 
g
#Figure 1: Comparing Feasibility of Different Designs
ggsave('output/all_design_permutations.tiff', g, width = 8.25, height = 5, units = 'in')